Repository navigation
feat: datacube shared-state for AnalysisInterferometer via curvature preloads - #566
Merged
Merged
Conversation
…preloads Implement the lensing consumer of the FactorGraphModel cross-Analysis shared-state mechanism for the interferometer datacube. - AnalysisInterferometer gains an opt-in `shared_preloads` flag and a `shared_state_from(instance)` that builds the channel-invariant curvature matrix once and returns it in a `PreloadsInterferometer`. - `log_likelihood_function(instance, shared=None)` and `fit_from(instance, preloads=None)` thread the shared object down through `FitInterferometer` and `TracerToInversion` into `inversion_from`, so each channel reuses the curvature matrix instead of rebuilding it. - `PreloadsInterferometer` is registered as a JAX pytree so the shared object threads through `jax.jit` (verified: shared-path likelihood matches numpy). Default behaviour is unchanged (`shared_preloads=False` / `shared=None`). Depends on the PyAutoArray Preloads API (must merge first). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Reuse the channel-invariant mapper across datacube channels, in addition to the curvature matrix, so the Delaunay triangulation is built once rather than per channel (~97% of the inversion-setup cost now shareable). - TracerToInversion.mapper_galaxy_dict short-circuits to a preloaded mapper when present (linear_obj_list and linear_obj_galaxy_dict both derive from it, so the mapper is not rebuilt anywhere on the hot path). - AnalysisInterferometer.shared_state_from builds one TracerToInversion and preloads both the curvature matrix and the mapper_galaxy_dict. - preloads ride as no_flatten config-aux on FitInterferometer (built inside the jit and consumed via attribute access, never crossing a jit boundary), so the PreloadsInterferometer pytree registration is removed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
Updated ( |
This was referenced Jun 8, 2026
Merged
3 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements the lensing consumer of the
FactorGraphModelcross-Analysisshared-state mechanism (PyAutoFit, shipped) for the interferometer datacube (PyAutoLens#565 / theanalysis_shared_stateepic). When a datacube is fit as NAnalysisInterferometerchannels sharing one lens model, the curvature matrixF = LᵀW̃L— the dominant per-channel inversion-setup cost — is identical for every channel. This builds it once on the lead factor and reuses it across every channel via the PyAutoArrayPreloadsAPI, targeting a large reduction in the inversion-setup cost for a many-channel cube.This is the curvature-sharing increment (the dominant ~86% of the shareable inversion-setup work). Mapper-sharing (the Delaunay rebuild) and the workspace scripts / profiling re-measure follow.
Depends on
Jammy2211/PyAutoArray(Preloads API) — that PR must merge first.API Changes
All additive and backward compatible — no removals or renames.
AnalysisInterferometer.__init__gains an opt-inshared_preloads=Falseflag.AnalysisInterferometer.shared_state_from(instance)— returns aPreloadsInterferometerwith the channel-invariantcurvature_matrix(orNonewhen not opted in / no inversion).AnalysisInterferometer.log_likelihood_functiongainsshared=None;fit_fromgainspreloads=None;FitInterferometerandTracerToInversiongainpreloads=None, threaded intoinversion_from.See full details below.
Test Plan
python -m pytest test_autolens/— full suite passes (331 tests)test_autolens/interferometer/model/test_analysis_interferometer.py:shared_state_fromyields aPreloadsInterferometerwhosecurvature_matrixis reused by the fit (identity) with unchangedfigure_of_merit;log_likelihood_function(shared=...)matches the unshared call; returnsNonewhen not opted in.PreloadsInterferometerthreads throughjax.jitand the shared-path likelihood matches numpy.Full API Changes (for automation & release notes)
Added
AnalysisInterferometer.shared_state_from(self, instance)— builds the sharedPreloadsInterferometeronce per evaluation (theautofit.Analysis.shared_state_fromhook for interferometer datacubes).Changed Signature (backward compatible — new optional args / flags)
AnalysisInterferometer.__init__(... , shared_preloads=False)AnalysisInterferometer.log_likelihood_function(self, instance, shared=None)AnalysisInterferometer.fit_from(self, instance, preloads=None)FitInterferometer.__init__(... , preloads=None)TracerToInversion.__init__(... , preloads=None)Changed Behaviour
shared_preloads=Trueand aFactorGraphModelforwards the sharedPreloadsInterferometer, each channel's inversion reuses the preloaded curvature matrix instead of rebuilding it. With the defaults, behaviour is unchanged.Migration
AnalysisInterferometer(..., shared_preloads=True)and combine them in aFactorGraphModel(which callsshared_state_fromand forwards the result).🤖 Generated with Claude Code